#export LD_PRELOAD=/usr/lib/libmpi.so python
from libatomism import *
from math import *
import array 
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib.mlab import griddata

def plotPES3d(kinop,repro,plt) :

   gcoors=kinop.getCoordinates()
   unit = gcoors.getUnits()
   system = kinop.getSystem()
   gcoorsg = msGeneralizedCoordinates.New(unit)
   q0 = msScalarVariable.New("Angstrom").set(1., 0., 20, 0.05, 0.05).setId("q0")
   q1 = msScalarVariable.New("Angstrom").set(1., 0., 20, 0.05, 0.05).setId("q1")
   gcoorsg.addVar(q0).addVar(q1)
 
   q0=gcoors.getVariable(0)
   q1=gcoors.getVariable(1)
   X = np.arange(q0.Min,q0.Max,q0.Dq)
   Y = np.arange(q1.Min,q1.Max,q1.Dq)
   scan = msSamplerLinear.New()
   scan.begin()
   Z = np.ones( (len(Y),len(X) ))   
   
   while( scan.getNextPoint( gcoors ) ):
       kinop.setDynamicDofs()
       system.computeCartCoordinates()
       vec = VectorOfDouble()
       vec.extend([system.separation(0,2)**(-2),system.separation(1,3)**(-2)])
       a =  repro.evaluate(vec)
       i = int( (q1.Value-q1.Min)/q1.Dq +0.5)
       j = int( (q0.Value-q0.Min)/q0.Dq +0.5)
       Z[i][j] = a
      
   X, Y = np.meshgrid(X, Y)
   surf = plt.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.coolwarm,
        linewidth=0, antialiased=False)
   plt.azim = 130
   plt.elev = 30
   plt.set_xlabel("reaction coordinate [Angstrom]",fontsize=20)
   plt.set_ylabel("torsional angle [degrees]",fontsize=20)
   plt.set_zlabel("energy [kcal/mol]",fontsize=20)
   plt.tick_params(labelsize=20)

def plotPES(epot,plt) :

   gcoors=epot.getCoordinates()
   unit = gcoors.getUnits()

   q0=gcoors.getVariable(0)
   q1=gcoors.getVariable(1)
   X = np.arange(q0.Min,q0.Max,q0.Dq)
   Y = np.arange(q1.Min,q1.Max,q1.Dq)
   scan = msSamplerLinear.New()
   scan.begin()
   Z  = np.ones( (len(Y),len(X) )) 
   
   while( scan.getNextPoint( gcoors ) ):
       Z[(q1.Value-q1.Min)/q1.Dq][(q0.Value-q0.Min)/q0.Dq] = epot.evaluate()
      
   X, Y = np.meshgrid(X, Y)
  
   levels=[19,21,23,24.2,26,28]
   CS = plt.contour(X, Y, Z,levels)
   for c in CS.collections:
     c.set_dashes([(0, (2.0, 2.0))])
   plt.axis((1.6,2.3,0,180))
   plt.set_xlabel("reaction coordinate [Angstrom]",fontsize=20)
   plt.set_ylabel("torsional angle [degrees]",fontsize=20)
   plt.tick_params(labelsize=20)

def plotRepro(kinop,repro,plt) :

   gcoors=kinop.getCoordinates()
   unit = gcoors.getUnits()
   system = kinop.getSystem()
   gcoorsg = msGeneralizedCoordinates.New(unit)
   q0 = msScalarVariable.New("Angstrom").set(1., 0., 20, 0.05, 0.05).setId("q0")
   q1 = msScalarVariable.New("Angstrom").set(1., 0., 20, 0.05, 0.05).setId("q1")
   gcoorsg.addVar(q0).addVar(q1)
 
   q0=gcoors.getVariable(0)
   q1=gcoors.getVariable(1)
   X = np.arange(q0.Min,q0.Max,q0.Dq)
   Y = np.arange(q1.Min,q1.Max,q1.Dq)
   scan = msSamplerLinear.New()
   scan.begin()
   Z = np.ones( (len(Y),len(X) ))   
   
   while( scan.getNextPoint( gcoors ) ):
       kinop.setDynamicDofs()
       system.computeCartCoordinates()
       vec = VectorOfDouble()
       vec.extend([system.separation(0,2)**(-2),system.separation(1,3)**(-2)])
       a =  repro.evaluate(vec)
       Z[(q1.Value-q1.Min)/q1.Dq][(q0.Value-q0.Min)/q0.Dq] = a

   X, Y = np.meshgrid(X, Y)
   levels=[19,21,23,24.2,26,28]
   CS = plt.contour(X, Y, Z,levels,linewidths=3.,colors='k')
   plt.axis((1.6,2.3,0,180))
   #im = plt.imshow(Z, cmap=cm.gray, extent=(1.6,2.3,0,180))
   plt.clabel(CS, inline=1, fontsize=16,backgroundcolor='m', color='b')

   with open('DataAI.txt') as f:
     w = f.readline().split()
     for line in f: 		
         array = [float(x) for x in line.split()]
         array.remove(array[len(array)-1])
         vec = VectorOfDouble()
         plt.plot(array[0],array[1],'*r', markersize=20)
   plt.set_xlabel("reaction coordinate [Angstrom]",fontsize=20)
   plt.set_ylabel("torsional angle [degrees]",fontsize=20)
   plt.tick_params(labelsize=20)

def histoCorrel(postPDFsample,i,j,plt):

# Create some random numbers
   x = []
   y = []
   param = postPDFsample.getUqParameters()
   postPDFsample.begin()
   k=0
   while(postPDFsample.getNextPoint()):
       x.append(param[i].Value)
       y.append(param[j].Value)    
       k=k+1
   # Estimate the 2D histogram
   nbins = 10
   H, xedges, yedges = np.histogram2d(x,y,bins=nbins)
   while(H.max()>100):
       nbins = nbins + 1 
       H, xedges, yedges = np.histogram2d(x,y,bins=nbins)
  
   # H needs to be rotated and flipped 
   H = np.rot90(H)
   H = np.flipud(H)
   z=[]
   postPDFsample.begin()
   while(postPDFsample.getNextPoint()):
       ix = min( len(H)*( param[i].Value - param[i].Min ) / ( param[i].Max - param[i].Min ) ,  len(H) -1 )
       iy = min( len(H)*( param[j].Value - param[j].Min ) / ( param[j].Max - param[j].Min ) ,  len(H) -1 )
       ix = max( 0, ix)
       iy=max(0,iy)
       z.append( H[ ix ] [ iy ]   )
      
   # Mask zeros
   Hmasked = np.ma.masked_where(H==0,H) # Mask pixels with a value of zero
   # Plot 2D histogram using pcolor
   # plt.pcolormesh(xedges,yedges,Hmasked)
   plt.scatter(x,y,z,c=z)
   plt.set_xlabel("param 0",fontsize=20)
   plt.set_ylabel("param 1",fontsize=20)
   plt.axis((1.8,3.0,3.4,4.4))
   plt.tick_params(labelsize=20)
